# -*- coding: utf-8 -*-
# time: 2025/3/26 08:52
# file: llm_qwen_interl.py
# author: hanson
"""
通过 HuggingFaceLLM 载入大模型，测试 qwen2.5
"""
import torch
from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline

from langchain_core.prompts import PromptTemplate, ChatPromptTemplate

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

# 2. 正确加载Qwen模型
model_id = r"E:\soft\model\qwen\Qwen\Qwen2___5-0___5B-Instruct"
# AutoTokenizer.from_pretrained用于加载预训练的tokenizer，负责将文本转换为模型可理解的数字形式。
tokenizer = AutoTokenizer.from_pretrained(
    pretrained_model_name_or_path=model_id,
    #model_id,
                                          trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,   # model_id,
    torch_dtype=torch.bfloat16,  # 设置模型精度
    device_map="auto",  # 设置模型在哪个设备上运行，这里设置为自动选择
    trust_remote_code=True  # 是否信任远程代码，如果模型是私有的，需要设置为False
)

# 创建文本生成pipeline
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=200,
    temperature=0.7
)

# 创建LangChain的LLM包装器
llm = HuggingFacePipeline(pipeline=pipe)

# 3. 修正prompt模板 - 使用ChatPromptTemplate
prompt_template = ChatPromptTemplate.from_template("""
您是一个设计用于查询文档来回答问题的代理。您可以使用文档检索工具。
并基于检索内容来回答问题您可能不查询文档就知道答案，但是您仍然应该查询文档来获得答案。
如果您从文档中找不到任何信息用于回答问题，则只需返回"抱歉，这个问题我还不知道。"作为答案。
如果有人提问等关于您的名字的问题，您就回答："我是超级牛逼哄哄的小天才助手"作为答案。

上下文：{context}
问题：{question}
""")

# 4. 创建正确的chain
chain = prompt_template | llm

# 5. 正确调用 - 需要传入字典参数
response = chain.invoke({
    "context": "无",  # 如果没有上下文可以设为"无"
    "question": "列出三种创造性的解决问题的技巧？"
})

# 6. 正确获取响应 - HuggingFacePipeline返回的是字符串
print(f"\n回答：\n{response}")